from settings import *
import geopandas as geopd
from geopandas.geoseries import *
from shapely.ops import linemerge
from shapely.geometry import LineString


class GeomFixFailException(Exception):
    def __init__(self, geom):
        super().__init__()
        self.geom = geom

    def __str__(self):
        return f"geom fix failed:{self.geom.wkt}"


def postgis_create_index(table_name, col, using=""):
    with conn:
        with conn.cursor() as cursor:
            sql_str = f"create index on {table_name} {using} ({col})"
            cursor.execute(sql_str)


def dtype_mapping(df):
    type_dict = {}
    for col, dtype in df.dtypes.items():
        dtype = str(dtype)
        col = col.lower()
        if dtype == "geometry":  # 只处理linestring
            geom = df.loc[0, col]
            if geom.geom_type == "LineString":
                if geom.has_z:
                    type_dict[col] = f"geometry(LINESTRINGZ, {SRID})"
                else:
                    type_dict[col] = f"geometry(LINESTRING, {SRID})"
        elif dtype == "object":
            type_dict[col] = "text"
        elif dtype == "float64":
            type_dict[col] = "float"
        elif dtype == "int64":
            type_dict[col] = "bigint"
    return type_dict


def create_table_by_df(df, table_name, drop_exist=True):
    # 根据df的column设计表
    exist = pd.read_sql(f"select count(*) from pg_class where relname = '{table_name}'", conn).values[0, 0]
    if exist and not drop_exist:
        return
    type_dict = dtype_mapping(df)
    sql_str = f"create table if not exists {table_name} ("
    for col, dtype in type_dict.items():
        sql_str += f"{col} {dtype},"
    sql_str = sql_str[:-1] + ")"  # 去掉最后一个逗号，加上反括号
    with conn:
        with conn.cursor() as cursor:
            if drop_exist:
                cursor.execute(f"drop table IF EXISTS {table_name}")
            cursor.execute(sql_str)
            conn.commit()


def postgis_create_backup(old_table, suffix):
    new_table = f"{old_table}_{suffix}"
    with conn:
        with conn.cursor() as cursor:
            sql_str = f"""
                            drop table if exists {new_table};
                            select * into {new_table} from {old_table};
                            alter table {new_table} add column if not exists {PK_COL} integer;
                            alter table {new_table} add primary key ({PK_COL})"""
            cursor.execute(sql_str)


def postgis_force2d(table_name):
    with conn.cursor() as cursor:
        old_geom_col = GEOM_COL + "_3d"
        sql_str = f"""
            begin;
            alter table {table_name} rename column {GEOM_COL} to {old_geom_col};
            alter table {table_name} add column {GEOM_COL} geometry(LINESTRING, {SRID});
            update {table_name} set {GEOM_COL} = st_force2d({old_geom_col});
            commit;
        """
        cursor.execute(sql_str, conn)


def geom_to_ewkt(geom, srid):
    return f"st_GeomFromEWKT('SRID={srid};{geom.wkt}')"


def postgis_add_geomcol(table_name, conn, EPGS, geomtype, print_text=True):
        cur = conn.cursor()
        sql = f"""SELECT AddGeometryColumn ('{table_name}','geom','{EPGS}','{geomtype}',2);"""
        try:
            cur.execute(sql)
            conn.commit()
            if print_text:
                print('add geom column')
        except:
            conn.commit()
            pass
